# -*- coding:utf-8 -*-
import re
import sqlite3
from prompt_toolkit.completion import Completer, Completion
from constants import SQL_KEYWORDS  # 修改为直接导入

class SQLCompleter(Completer):
    def __init__(self, cli):
        self.cli = cli
        self.keywords = SQL_KEYWORDS
        self.table_pattern = re.compile(r'from\s+([a-zA-Z_][a-zA-Z0-9_]*)', re.IGNORECASE)
        
    def get_table_columns(self, table_name):
        try:
            self.cli.cursor.execute(f"PRAGMA table_info({table_name})")
            columns = self.cli.cursor.fetchall()
            return [(row[1], row[2]) for row in columns]
        except sqlite3.Error:
            return []

    def get_completions(self, document, complete_event):
        word_before_cursor = document.get_word_before_cursor()
        text = document.text_before_cursor.lower()

        if 'select' in text:
            table_match = self.table_pattern.search(text)
            if table_match:
                table_name = table_match.group(1)
                columns = self.get_table_columns(table_name)
                if ('from' not in text[text.rfind('select'):] or 
                    'where' in text or 
                    'order by' in text):
                    for column, col_type in columns:
                        if column.lower().startswith(word_before_cursor.lower()):
                            yield Completion(column, 
                                          start_position=-len(word_before_cursor),
                                          display_meta=f"({col_type})")
                    return

        if 'insert into' in text:
            words = text.split()
            if len(words) >= 3:
                table_name = words[2]
                columns = self.get_table_columns(table_name)
                if '(' in text and ')' not in text:
                    for column, col_type in columns:
                        if column.lower().startswith(word_before_cursor.lower()):
                            yield Completion(column, 
                                          start_position=-len(word_before_cursor),
                                          display_meta=f"({col_type})")
                    return

        if 'update' in text and 'set' in text:
            words = text.split()
            if len(words) >= 2:
                table_name = words[1]
                columns = self.get_table_columns(table_name)
                for column, col_type in columns:
                    if column.lower().startswith(word_before_cursor.lower()):
                        yield Completion(column, 
                                      start_position=-len(word_before_cursor),
                                      display_meta=f"({col_type})")
                return

        words = self.keywords + self.cli.get_table_names()
        for word in words:
            if word.lower().startswith(word_before_cursor.lower()):
                yield Completion(word, start_position=-len(word_before_cursor))